package device

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"net"
	"time"

	"golang.org/x/crypto/chacha20poly1305"
)

const RejectAfterMessages = (1 << 64) - (1 << 13) - 1

func (device *Device) routineHandshake() {
	device.Wg.Add(1)
	defer device.Wg.Done()
	for {
		select {
		case <-device.Done:
			return
		case packet := <-device.queue.handshake:
			{
				msgType := packet.msgType
				switch msgType {
				case 1:
					// 消费握手请求消息
					if len(packet.packet) != 148 {
						fmt.Println("握手请求消息长度不正确")
						goto skip
					}
					if !device.cookieChecker.checkMAC1(packet.packet) {
						fmt.Println("握手消息MAC1校验失败")
						goto skip
					}

					//todo  当任务队列过载时需要校验mac2
					var msg MessageInitiation
					reader := bytes.NewReader(packet.packet)
					err := binary.Read(reader, binary.LittleEndian, &msg)
					if err != nil {
						fmt.Printf("Failed to decode initiation message")
						goto skip
					}

					// consume initiation

					peer := device.consumeMessageInitiation(&msg)
					if peer == nil {
						fmt.Printf("Received invalid initiation message ")
						goto skip
					}

					peer.timersAnyAuthenticatedPacketTraversal()
					peer.timersAnyAuthenticatedPacketReceived()
					peer.updateEndpoint(packet.addr)
					peer.sendHandshakeResponse()
				case 2:
					// 消费握手响应消息
					if len(packet.packet) != 92 {
						goto skip
					}
					if !device.cookieChecker.checkMAC1(packet.packet) {
						goto skip
					}
					// unmarshal

					var msg MessageResponse
					reader := bytes.NewReader(packet.packet)
					err := binary.Read(reader, binary.LittleEndian, &msg)
					if err != nil {
						fmt.Println("Failed to decode response message")
						goto skip
					}
					// consume response
					peer := device.ConsumeMessageResponse(&msg)
					if peer == nil {
						goto skip
					}
					peer.updateEndpoint(packet.addr)
					// derive keypair
					err = peer.beginSymmetricSession()
					peer.timersAnyAuthenticatedPacketTraversal()
					peer.timersHandshakeComplete()
					peer.timersAnyAuthenticatedPacketReceived()

					if err != nil {
						fmt.Println("Failed to derive keypair")
						goto skip
					}
					fmt.Println("握手成功")
				case 3:
					// 消费cookie信息
					var reply MessageCookieReply
					reader := bytes.NewReader(packet.packet)
					err := binary.Read(reader, binary.LittleEndian, &reply)
					if err != nil {
						fmt.Println("Failed to decode cookie reply")
						goto skip
					}
					// lookup peer from index
					entry, ok := device.indexTable.lookup(reply.Receiver)
					if !ok {
						goto skip
					}
					peer := entry.peer
					if !peer.cookieGenerator.consumeReply(&reply) {
						fmt.Println("Could not decrypt invalid cookie response")
					}
				default:
					fmt.Println("收到未知握手消息")
				}
				goto skip
			}
		skip:
			// 回收到对象池
			device.pool.Put(packet.buffer)
		}
	}
}

func (device *Device) routineTansportData() {
	device.Wg.Add(1)
	defer device.Wg.Done()
	var nonce [chacha20poly1305.NonceSize]byte
	for {
		select {
		case <-device.Done:
			return
		case packet := <-device.queue.recv:
			{
				// check size , header:16 + 加密tagsize:16
				if len(packet.packet) < 32 {
					goto skip
				}

				// lookup key pair
				receiver := binary.LittleEndian.Uint32(
					packet.packet[4:8],
				)
				value, ok := device.indexTable.lookup(receiver)
				if !ok {
					fmt.Println("未找到对应的密钥对")
					goto skip
				}
				keypair := value.keypair
				if keypair == nil {
					fmt.Println("密钥对为空")
					goto skip
				}

				// check keypair expiry
				if keypair.created.Add(time.Second * 180).Before(time.Now()) {
					fmt.Println("密钥对已过期")
					goto skip
				}
				counter := binary.LittleEndian.Uint64(packet.packet[8:16])
				content := packet.packet[16:]
				var err error
				// copy counter to nonce
				binary.LittleEndian.PutUint64(nonce[4:12], counter)
				packet.packet, err = keypair.receive.Open(
					content[:0],
					nonce[:],
					content,
					nil,
				)
				if err != nil {
					fmt.Println("解密失败")
					goto skip
				}

				if !keypair.replayFilter.ValidateCounter(counter, RejectAfterMessages) {
					goto skip
				}
				peer := value.peer
				if peer.receivedWithKeypair(keypair) {
					peer.timersHandshakeComplete()
				}
				srcIP, dstIP := parseIPHeader(packet.packet)

				if device.AllowedIPs.find(srcIP) != peer {
					fmt.Printf("不支持的源地址 %s -> %s\n", srcIP.String(), dstIP.String())
					goto skip
				}
				device.Tun.Write(packet.packet)
				peer.timersAnyAuthenticatedPacketTraversal()
				peer.timersAnyAuthenticatedPacketReceived()
				peer.keepKeyFreshReceiving()
				peer.updateEndpoint(packet.addr)
				goto skip
			}
		skip:
			// 回收到对象池
			device.pool.Put(packet.buffer)
		}

	}
}

func parseIPHeader(packet []byte) (srcIP, dstIP net.IP) {
	if len(packet) < 20 {
		return nil, nil
	}
	srcIP = net.IP(packet[12:16])
	dstIP = net.IP(packet[16:20])
	return
}

func (peer *Peer) keepKeyFreshReceiving() {
	if peer.timers.sentLastMinuteHandshake.Load() {
		return
	}
	keypair := peer.keypairs.Current()
	if keypair != nil && keypair.isInitiator && time.Since(keypair.created) > (RejectAfterTime-KeepaliveTimeout-RekeyTimeout) {
		peer.timers.sentLastMinuteHandshake.Store(true)
		peer.sendHandshakeInitiation(false)
	}
}
